Rmalliu/feat/multi share mps - #12

Summary

Summary

Enable pods on MPS-enabled GPU clusters to request more than one shared replica and, when they request all replicas, get the full GPU hardware. Two coordinated changes in this PR:

1. Make failRequestsGreaterThanOne authoritative for MPS (commit 793753a1). Today the flag is overridden to true and then ignored, so MPS resource requests > 1 are always rejected regardless of config. This change honors the configured value the same way time-slicing already does. The default remains
   true (reject multi-replica requests) for backward compatibility — clusters opt in by explicitly setting sharing.mps.failRequestsGreaterThanOne: false in the device-plugin config.
2. Grant full GPU resources when a container requests every replica advertised on the node (commit fcdd6c77). MPS control daemon's server-side caps are raised to full per-device memory and 100% threads, and the Allocate() response injects CUDA_MPS_ACTIVE_THREAD_PERCENTAGE and
   CUDA_MPS_PINNED_DEVICE_MEM_LIMIT at 1/replicas for any non-full-node grant — preserving today's behavior for partial requests, while letting all-replica grants ride unrestricted daemon defaults.

Motivation

The MPS control daemon sets per-client per-device caps at startup:

• set_default_device_pinned_mem_limit <index> <totalMem/replicas>
• set_default_active_thread_percentage <100/replicas>

Per NVIDIA's MPS docs (appendix on CUDA_MPS_PINNED_DEVICE_MEM_LIMIT):

"The new limit will only further constrain the limit set by the control daemon … setting this environment variable to a value greater than value observed by the server will not set the limit to the higher value and thus will be ineffective."

So the daemon default is a hard cap — client env vars can request less, never more. The result: even a container that's been allocated every replica ID on a node is still throttled to 1/replicas per physical GPU. nvidia-smi correctly shows all GPUs because the device nodes are mounted; the MPS server
silently caps allocations.

How it works

Two coordinated touchpoints:

1. cmd/mps-control-daemon/mps/daemon.go — daemon defaults raised:

• perDevicePinnedDeviceMemoryLimits() → full device.TotalMemory per physical index
• activeThreadPercentage() → 100

These were the floor for every client; raising them lets the client env vars (which can only clamp down, never up) become the effective cap.

2. internal/plugin/server.go + internal/plugin/mps.go — Allocate response:

• updateResponseForMPS and mpsOptions.updateReponse now take grantedCount int (= len(req.DevicesIDs)).
• If grantedCount >= total advertised replicas, no env vars are injected; the container inherits the daemon's full-hardware defaults.
• Otherwise, the response injects:
  • CUDA_MPS_ACTIVE_THREAD_PERCENTAGE = 100/replicas
  • CUDA_MPS_PINNED_DEVICE_MEM_LIMIT = "<idx>=<totalMem/replicas>M,<idx>=<totalMem/replicas>M,..."

Both per-physical-GPU values come from m.daemon.Devices() — already in scope, no new plumbing of device IDs needed.

Why this is safe (no oversubscription)

The full-grant branch (grantedCount >= total) only triggers when kubelet's device-plugin accounting has handed this single container every replica ID the plugin advertises. Kubelet's accounting guarantees:

• Each replica ID is held by ≤1 pod at any instant.
• The all-replicas request can only be admitted when every ID is free.
• No other pod can be admitted (for this resource) while this pod runs.

So at the moment unrestricted defaults apply, exactly one MPS client exists on the node. NVIDIA's documented hard cap of 60 client CUDA contexts per device (Volta+) is never the active constraint — kubelet refuses oversubscription long before MPS would.

The "mixed allocation" footgun (a pod gets all replicas of GPU-A + some of GPU-B on a multi-GPU node) cannot satisfy grantedCount == total by construction; it's all-or-nothing.

Operational notes

• Daemon restart is required for the new defaults to take effect. set_default_* is applied once at MPS daemon startup; existing clients keep their old caps. Rolling the daemonset disrupts in-flight MPS clients — coordinate with workloads.
• HPC-V on MPS nodes. The MPS device-plugin injects a /dev/shm bind mount into every MPS client. cwtest's per-node test pods currently also create an emptyDir{medium:Memory} at /dev/shm, stacking the two and triggering a gamble [7D07] sys_disk_mounts false negative. This PR doesn't touch that —
  follow-up changes in coreweave/getnodes (skip the emptyDir on MPS nodes) and coreweave/gamble (env-var-aware /dev/shm source check) are needed to keep HPC-V passing.
• Multi-resource setups. If a node advertises more than one MPS-managed resource type, each gets its own control daemon; this change applies to all of them automatically (it's in Daemon.Start()).

What this does NOT cover

• Per-physical-GPU proportional grants for mixed/fractional requests. Anything that isn't a full-node grant gets the same 1/replicas per-device cap as before, regardless of how many replicas it owns of any specific physical GPU. The framework's there if it's ever wanted (group req.DevicesIDs by UUID and
  count per-physical), but not needed for the current goal.
• Per-PID set_active_thread_percentage <PID> via the control daemon (alternative to env vars). More complex (host PID plumbing, fork/exec doesn't propagate per-PID settings) and not needed for the all-or-nothing case.
• MIG. Already rejected for MPS in internal/plugin/mps.go.

Test plan

• Build clean: go build ./...
• Existing tests pass: go test ./cmd/mps-control-daemon/... ./internal/plugin/...
• Deploy to mps-lab (g884dd0, 8× RTX PRO 6000 Blackwell, MPS replicas=10).
• Restart the nvidia-device-plugin-mps-control-daemon daemonset; confirm logs show set_default_active_thread_percentage 100 and set_default_device_pinned_mem_limit <idx> <fullMem>M for each GPU.
• Schedule a pod requesting nvidia.com/gpu: "10" (one full physical GPU's worth of replicas). Inside the pod:
  • Verify CUDA_MPS_ACTIVE_THREAD_PERCENTAGE and CUDA_MPS_PINNED_DEVICE_MEM_LIMIT env vars are not set (full-node-per-GPU branch path), OR the equivalent 100% / full-memory values depending on how the resource was allocated.
  • Run a CUDA program that allocates more than totalMem/10 pinned memory on the assigned GPU; confirm it succeeds (previously would cudaErrorOutOfMemory).
  • Run a kernel sized to the full SM count; observe throughput close to a dedicated-GPU baseline.
• Schedule a pod requesting nvidia.com/gpu: "1" (partial). Confirm:
  • CUDA_MPS_ACTIVE_THREAD_PERCENTAGE=10 is set.
  • CUDA_MPS_PINNED_DEVICE_MEM_LIMIT contains entries for each physical GPU at totalMem/10M.
  • A CUDA program allocating more than totalMem/10 fails with cudaErrorOutOfMemory — i.e. the cap is still enforced for partial grants.
• Schedule a pod requesting all replicas on the node (nvidia.com/gpu: "80" for the lab node). Confirm no env vars injected and CUDA can use the full hardware on every GPU.

References

• NVIDIA MPS — CUDA_MPS_PINNED_DEVICE_MEM_LIMIT / CUDA_MPS_ACTIVE_THREAD_PERCENTAGE semantics: https://docs.nvidia.com/deploy/mps/appendix-tools-and-interface-reference.html
• NVIDIA MPS — 60-client-per-device cap and connection-refusal behavior: https://docs.nvidia.com/deploy/mps/when-to-use-mps.html

🤖 Generated with Claude Code